"""
时间序列预测模型模块
"""

import numpy as np
import pandas as pd
from typing import Dict, List, Tuple, Optional, Union
from statsmodels.tsa.arima.model import ARIMA
from statsmodels.tsa.holtwinters import ExponentialSmoothing
from statsmodels.tsa.vector_ar.var_model import VAR
from .ilr_transformation import ILRTransformation
import warnings
warnings.filterwarnings('ignore')


class BaseForecaster:
    """基础预测器类"""
    
    def __init__(self, **kwargs):
        """初始化预测器"""
        self.model = None
        self.fitted = False
        self.params = kwargs
        
    def fit(self, y: np.ndarray) -> 'BaseForecaster':
        """拟合模型"""
        raise NotImplementedError
        
    def forecast(self, horizon: int) -> np.ndarray:
        """预测"""
        raise NotImplementedError
        
    def fit_forecast(self, y: np.ndarray, horizon: int) -> np.ndarray:
        """拟合并预测"""
        return self.fit(y).forecast(horizon)


class ARIMAForecaster(BaseForecaster):
    """ARIMA预测器"""
    
    def __init__(self, order: Tuple[int, int, int] = None, **kwargs):
        """
        初始化ARIMA预测器
        
        Args:
            order: ARIMA阶数 (p, d, q)
        """
        super().__init__(**kwargs)
        self.order = order or (1, 1, 1)
        
    def fit(self, y: np.ndarray) -> 'ARIMAForecaster':
        """拟合ARIMA模型"""
        try:
            # 自动确定最优阶数
            if self.order == (1, 1, 1):
                from statsmodels.tsa.arima.model import ARIMA
                from statsmodels.tsa.stattools import adfuller
                
                # 检查平稳性
                adf_result = adfuller(y)
                if adf_result[1] > 0.05:
                    d = 1
                else:
                    d = 0
                    
                # 简单的阶数选择
                p = min(2, len(y) // 10)
                q = min(2, len(y) // 10)
                self.order = (p, d, q)
            
            self.model = ARIMA(y, order=self.order)
            self.fitted_model = self.model.fit()
            self.fitted = True
            
        except Exception as e:
            # 如果自动选择失败，使用默认阶数
            self.order = (1, 1, 1)
            self.model = ARIMA(y, order=self.order)
            self.fitted_model = self.model.fit()
            self.fitted = True
            
        return self
        
    def forecast(self, horizon: int) -> np.ndarray:
        """ARIMA预测"""
        if not self.fitted:
            raise ValueError("请先调用fit方法")
            
        forecast_result = self.fitted_model.forecast(steps=horizon)
        return forecast_result.values


class ETSForecaster(BaseForecaster):
    """指数平滑预测器"""
    
    def __init__(self, trend: str = 'add', seasonal: str = None, **kwargs):
        """
        初始化ETS预测器
        
        Args:
            trend: 趋势类型 ('add', 'mul', None)
            seasonal: 季节性类型 ('add', 'mul', None)
        """
        super().__init__(**kwargs)
        self.trend = trend
        self.seasonal = seasonal
        
    def fit(self, y: np.ndarray) -> 'ETSForecaster':
        """拟合ETS模型"""
        try:
            self.model = ExponentialSmoothing(
                y,
                trend=self.trend,
                seasonal=self.seasonal,
                **self.params
            )
            self.fitted_model = self.model.fit()
            self.fitted = True
            
        except Exception as e:
            # 如果拟合失败，使用简单指数平滑
            self.model = ExponentialSmoothing(y)
            self.fitted_model = self.model.fit()
            self.fitted = True
            
        return self
        
    def forecast(self, horizon: int) -> np.ndarray:
        """ETS预测"""
        if not self.fitted:
            raise ValueError("请先调用fit方法")
            
        forecast_result = self.fitted_model.forecast(steps=horizon)
        return forecast_result.values


class THETAForecaster(BaseForecaster):
    """THETA预测器"""
    
    def __init__(self, **kwargs):
        """初始化THETA预测器"""
        super().__init__(**kwargs)
        
    def fit(self, y: np.ndarray) -> 'THETAForecaster':
        """拟合THETA模型"""
        # THETA方法：对去趋势数据进行简单指数平滑
        self.original_data = y.copy()
        
        # 计算线性趋势
        n = len(y)
        x = np.arange(n)
        slope, intercept = np.polyfit(x, y, 1)
        
        # 去趋势
        self.trend = slope * x + intercept
        self.detrended = y - self.trend
        
        # 对去趋势数据拟合简单指数平滑
        self.ets_model = ExponentialSmoothing(self.detrended)
        self.fitted_model = self.ets_model.fit()
        self.fitted = True
        
        return self
        
    def forecast(self, horizon: int) -> np.ndarray:
        """THETA预测"""
        if not self.fitted:
            raise ValueError("请先调用fit方法")
            
        # 预测去趋势部分
        detrended_forecast = self.fitted_model.forecast(steps=horizon)
        
        # 添加趋势
        n = len(self.original_data)
        future_x = np.arange(n, n + horizon)
        future_trend = self.trend[0] + (future_x - n) * (self.trend[1] - self.trend[0])
        
        return detrended_forecast.values + future_trend


class EnsembleForecaster(BaseForecaster):
    """集成预测器"""
    
    def __init__(self, forecasters: List[BaseForecaster] = None, weights: List[float] = None):
        """
        初始化集成预测器
        
        Args:
            forecasters: 预测器列表
            weights: 权重列表
        """
        super().__init__()
        self.forecasters = forecasters or [
            ARIMAForecaster(),
            ETSForecaster(),
            THETAForecaster()
        ]
        self.weights = weights or [1/len(self.forecasters)] * len(self.forecasters)
        
    def fit(self, y: np.ndarray) -> 'EnsembleForecaster':
        """拟合所有预测器"""
        self.fitted_models = []
        for forecaster in self.forecasters:
            fitted_model = forecaster.fit(y.copy())
            self.fitted_models.append(fitted_model)
        self.fitted = True
        return self
        
    def forecast(self, horizon: int) -> np.ndarray:
        """集成预测"""
        if not self.fitted:
            raise ValueError("请先调用fit方法")
            
        forecasts = []
        for model in self.fitted_models:
            forecast = model.forecast(horizon)
            forecasts.append(forecast)
            
        # 加权平均
        weighted_forecast = np.zeros(horizon)
        for i, (forecast, weight) in enumerate(zip(forecasts, self.weights)):
            weighted_forecast += weight * forecast
            
        return weighted_forecast


class ForecastingModels:
    """预测模型管理器"""
    
    def __init__(self, config: Dict = None):
        """
        初始化预测模型管理器
        
        Args:
            config: 配置字典
        """
        self.config = config or {}
        self.models = {}
        
    def create_model(self, model_name: str, **kwargs) -> BaseForecaster:
        """
        创建预测模型
        
        Args:
            model_name: 模型名称
            **kwargs: 模型参数
            
        Returns:
            预测器实例
        """
        if model_name.upper() == 'ARIMA':
            return ARIMAForecaster(**kwargs)
        elif model_name.upper() == 'ETS':
            return ETSForecaster(**kwargs)
        elif model_name.upper() == 'THETA':
            return THETAForecaster(**kwargs)
        elif model_name.upper() == 'ENSEMBLE':
            return EnsembleForecaster(**kwargs)
        else:
            raise ValueError(f"不支持的模型类型: {model_name}")
    
    def fit_models(self, data: pd.DataFrame, target_columns: List[str] = None) -> Dict[str, BaseForecaster]:
        """
        拟合多个模型
        
        Args:
            data: 时间序列数据框
            target_columns: 目标列名列表
            
        Returns:
            拟合后的模型字典
        """
        if target_columns is None:
            target_columns = [col for col in data.columns if col != 'Year']
            
        fitted_models = {}
        
        for column in target_columns:
            if column in data.columns:
                y = data[column].values
                
                # 创建集成模型
                ensemble_model = EnsembleForecaster()
                fitted_model = ensemble_model.fit(y)
                fitted_models[column] = fitted_model
                
        return fitted_models
    
    def forecast_all(self, models: Dict[str, BaseForecaster], horizon: int) -> pd.DataFrame:
        """
        对所有模型进行预测
        
        Args:
            models: 拟合后的模型字典
            horizon: 预测步长
            
        Returns:
            预测结果数据框
        """
        forecasts = {}
        
        for column, model in models.items():
            try:
                forecast = model.forecast(horizon)
                forecasts[column] = forecast
            except Exception as e:
                print(f"预测列 {column} 失败: {e}")
                forecasts[column] = np.full(horizon, np.nan)
                
        return pd.DataFrame(forecasts)
    
    def cross_validate(self, data: pd.DataFrame, target_columns: List[str], 
                      horizon: int, cv_splits: int = 5) -> Dict[str, float]:
        """
        交叉验证
        
        Args:
            data: 时间序列数据
            target_columns: 目标列名
            horizon: 预测步长
            cv_splits: 交叉验证折数
            
        Returns:
            验证结果字典
        """
        from sklearn.metrics import mean_squared_error, mean_absolute_error
        
        results = {}
        
        for column in target_columns:
            if column not in data.columns:
                continue
                
            y = data[column].values
            n = len(y)
            
            mse_scores = []
            mae_scores = []
            
            for i in range(cv_splits):
                # 划分训练集和测试集
                split_point = n - horizon - i * horizon
                if split_point < horizon:
                    break
                    
                train_data = y[:split_point]
                test_data = y[split_point:split_point + horizon]
                
                # 拟合模型
                model = EnsembleForecaster()
                fitted_model = model.fit(train_data)
                
                # 预测
                forecast = fitted_model.forecast(horizon)
                
                # 计算误差
                mse = mean_squared_error(test_data, forecast)
                mae = mean_absolute_error(test_data, forecast)
                
                mse_scores.append(mse)
                mae_scores.append(mae)
                
            if mse_scores:
                results[f"{column}_MSE"] = np.mean(mse_scores)
                results[f"{column}_MAE"] = np.mean(mae_scores)
                
        return results 


class VARForecaster(BaseForecaster):
    """VAR预测器"""
    
    def __init__(self, maxlags: int = None, **kwargs):
        """
        初始化VAR预测器
        
        Args:
            maxlags: 最大滞后阶数
        """
        super().__init__(**kwargs)
        self.maxlags = maxlags
        
    def fit(self, y: np.ndarray) -> 'VARForecaster':
        """拟合VAR模型"""
        try:
            # 如果是单变量数据，转换为二维
            if y.ndim == 1:
                y = y.reshape(-1, 1)
                
            # 检查变量数量
            if y.shape[1] < 2:
                raise ValueError(f"VAR模型需要至少2个变量，当前只有{y.shape[1]}个变量")
                
            # 创建DataFrame
            df = pd.DataFrame(y)
            
            # 自动选择最优滞后阶数
            if self.maxlags is None:
                self.maxlags = min(10, len(y) // 10)
                
            # 拟合VAR模型
            self.model = VAR(df)
            self.fitted_model = self.model.fit(maxlags=self.maxlags, ic='aic')
            self.fitted = True
            
        except Exception as e:
            # 如果自动选择失败，使用默认滞后阶数
            if 'df' in locals():
                self.maxlags = 1
                self.model = VAR(df)
                self.fitted_model = self.model.fit(maxlags=self.maxlags)
                self.fitted = True
            else:
                raise e
            
        return self
        
    def forecast(self, horizon: int) -> np.ndarray:
        """VAR预测"""
        if not self.fitted:
            raise ValueError("请先调用fit方法")
            
        forecast_result = self.fitted_model.forecast(
            self.fitted_model.endog, steps=horizon
        )
        # forecast_result 已经是 numpy 数组，不需要 .values
        return forecast_result


def forecast_ilr_var(ilr_data: np.ndarray, horizon: int, 
                    transformer: 'ILRTransformation' = None,
                    **var_kwargs) -> Tuple[np.ndarray, 'VARForecaster']:
    """
    使用VAR模型对ILR变换后的数据进行预测
    
    Args:
        ilr_data: ILR变换后的数据矩阵 (n_samples, n_components-1)
        horizon: 预测步数
        transformer: ILR变换器（用于逆变换）
        **var_kwargs: VAR模型的其他参数
        
    Returns:
        (预测的比例数据, VAR预测器)
    """
    print(f"开始VAR预测，数据形状: {ilr_data.shape}, 预测步数: {horizon}")
    
    # 检查变量数量，决定使用VAR还是ARIMA
    n_variables = ilr_data.shape[1] if ilr_data.ndim > 1 else 1
    
    if n_variables >= 2:
        # 变量数量大于等于2，使用VAR模型
        print(f"变量数量: {n_variables} >= 2，使用VAR模型")
        
        # 1. 训练VAR模型
        print("1. 训练VAR模型...")
        var_forecaster = VARForecaster(**var_kwargs)
        var_forecaster.fit(ilr_data)
        print("✓ VAR模型训练完成")
        
        # 2. 预测未来N步
        print(f"2. 预测未来{horizon}步...")
        ilr_forecast = var_forecaster.forecast(horizon)
        print(f"✓ ILR预测完成，预测结果形状: {ilr_forecast.shape}")
        
    else:
        # 变量数量小于2，使用ARIMA模型作为备用方案
        print(f"变量数量: {n_variables} < 2，切换到ARIMA模型作为备用方案")
        
        # 确保数据是二维的
        if ilr_data.ndim == 1:
            ilr_data = ilr_data.reshape(-1, 1)
        
        # 1. 使用ARIMA模型训练
        print("1. 训练ARIMA模型...")
        try:
            from pmdarima import auto_arima
            # 对第一个（也是唯一的）变量进行ARIMA预测
            series = ilr_data[:, 0]
            arima_model = auto_arima(
                series,
                start_p=0, start_q=0,
                max_p=5, max_q=5, max_d=2,
                seasonal=False,
                stepwise=True,
                suppress_warnings=True,
                error_action='ignore',
                trace=False
            )
            print(f"✓ ARIMA模型训练完成: {arima_model}")
            
            # 2. 预测未来N步
            print(f"2. 预测未来{horizon}步...")
            ilr_forecast = arima_model.predict(n_periods=horizon).reshape(-1, 1)
            print(f"✓ ARIMA预测完成，预测结果形状: {ilr_forecast.shape}")
            
            # 创建一个虚拟的VAR预测器来保持接口一致
            class DummyVARForecaster:
                def __init__(self, arima_model):
                    self.arima_model = arima_model
                    self.fitted = True
                
                def forecast(self, horizon):
                    return self.arima_model.predict(n_periods=horizon).reshape(-1, 1)
            
            var_forecaster = DummyVARForecaster(arima_model)
            
        except Exception as e:
            print(f"✗ ARIMA模型训练失败: {e}")
            raise e
    
    # 3. 对预测结果进行ILR逆变换
    print("3. 执行ILR逆变换...")
    if transformer is not None:
        proportions_forecast = transformer.inverse_transform(ilr_forecast)
        print(f"✓ ILR逆变换完成，比例预测结果形状: {proportions_forecast.shape}")
    else:
        print("⚠ 警告: 未提供ILR变换器，跳过逆变换步骤")
        proportions_forecast = ilr_forecast
    
    # 4. 验证比例数据的有效性
    if transformer is not None:
        row_sums = np.sum(proportions_forecast, axis=1)
        print(f"比例数据行和范围: [{row_sums.min():.6f}, {row_sums.max():.6f}]")
        
        # 确保比例和为1
        if not np.allclose(row_sums, 1.0, atol=1e-6):
            print("⚠ 警告: 比例数据行和不为1，进行标准化")
            proportions_forecast = proportions_forecast / row_sums.reshape(-1, 1)
    
    print("✓ VAR-ILR预测完成")
    
    return proportions_forecast, var_forecaster


def forecast_ilr_var_by_region(ilr_results: Dict[str, Dict], 
                             horizon: int,
                             **var_kwargs) -> Dict[str, Tuple[np.ndarray, 'VARForecaster']]:
    """
    按区域使用VAR模型对ILR变换后的数据进行预测
    
    Args:
        ilr_results: 各区域ILR变换结果字典
        horizon: 预测步数
        **var_kwargs: VAR模型的其他参数
        
    Returns:
        各区域的预测结果字典
    """
    print(f"开始按区域VAR预测，预测步数: {horizon}")
    
    forecast_results = {}
    
    for region, result in ilr_results.items():
        print(f"\n处理区域: {region}")
        
        try:
            ilr_coords = result['ilr_coords']
            transformer = result['transformer']
            
            # 执行VAR预测
            proportions_forecast, var_forecaster = forecast_ilr_var(
                ilr_coords, horizon, transformer, **var_kwargs
            )
            
            forecast_results[region] = {
                'proportions_forecast': proportions_forecast,
                'var_forecaster': var_forecaster,
                'columns': result['columns']
            }
            
            print(f"✓ 区域 {region} 预测完成")
            
        except Exception as e:
            print(f"✗ 区域 {region} 预测失败: {e}")
            continue
    
    print(f"\n✓ 所有区域预测完成，成功预测 {len(forecast_results)} 个区域")
    
    return forecast_results 